#include <dcmtk/dcmdata/dcfilefo.h>
#include <dcmtk/dcmdata/dcmetinf.h>
#include <dcmtk/dcmdata/dcdeftag.h>
#include <png.h>
#include <iostream>

void save_png_gray8(uint8_t *pixel_data, int rows)
{
    // libpng struct initialize
    auto png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
    if (png_ptr == nullptr)
    {
        std::cerr << "Failed to png_create_write_struct" << std::endl;
        png_destroy_write_struct(&png_ptr, nullptr);
        return;
    }

    auto info_ptr = png_create_info_struct(png_ptr);
    if (info_ptr == nullptr)
    {
        std::cerr << "Failed to png_create_info_struct" << std::endl;
        png_destroy_write_struct(&png_ptr, &info_ptr);
        return;
    }

    if (setjmp(png_jmpbuf(png_ptr)))
    {
        std::cerr << "Failed to setjmp" << std::endl;
        png_destroy_write_struct(&png_ptr, &info_ptr);
        return;
    }

    // assoc with file pointer, and write to file.
    FILE *fp = fopen("gray8.png", "wb");
    png_init_io(png_ptr, fp);
    png_set_filter(png_ptr, 0, PNG_FILTER_NONE);
    png_set_IHDR(png_ptr, info_ptr, 512, 512, 8, PNG_COLOR_TYPE_GRAY,
                 PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

    png_set_packing(png_ptr);
    png_write_info(png_ptr, info_ptr);

    uint8_t **raw_data = new uint8_t *[rows];
    for (int i = 0; i < 512; ++i)
        raw_data[i] = pixel_data + (i * 512);

    png_set_rows(png_ptr, info_ptr, raw_data);
    png_write_png(png_ptr, info_ptr, PNG_TRANSFORM_IDENTITY, NULL);
    png_write_end(png_ptr, info_ptr);
    fclose(fp);

    png_destroy_write_struct(&png_ptr, &info_ptr);
}

void save_png_gray16(uint16_t *pixel_data, int rows)
{
    // libpng struct initialize
    auto png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
    if (png_ptr == nullptr)
    {
        std::cerr << "Failed to png_create_write_struct" << std::endl;
        png_destroy_write_struct(&png_ptr, nullptr);
        return;
    }

    auto info_ptr = png_create_info_struct(png_ptr);
    if (info_ptr == nullptr)
    {
        std::cerr << "Failed to png_create_info_struct" << std::endl;
        png_destroy_write_struct(&png_ptr, &info_ptr);
        return;
    }

    if (setjmp(png_jmpbuf(png_ptr)))
    {
        std::cerr << "Failed to setjmp" << std::endl;
        png_destroy_write_struct(&png_ptr, &info_ptr);
        return;
    }

    // assoc with file pointer, and write to file.
    FILE *fp = fopen("gray16.png", "wb");
    png_init_io(png_ptr, fp);
    png_set_filter(png_ptr, 0, PNG_FILTER_NONE);
    png_set_IHDR(png_ptr, info_ptr, 512, 512, 16, PNG_COLOR_TYPE_GRAY,
                 PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

    png_set_packing(png_ptr);
    png_write_info(png_ptr, info_ptr);

    uint16_t **raw_data = new uint16_t *[rows];
    for (int i = 0; i < 512; ++i)
        raw_data[i] = pixel_data + (i * 512);

    png_set_rows(png_ptr, info_ptr, reinterpret_cast<uint8_t **>(raw_data));
    png_write_png(png_ptr, info_ptr, PNG_TRANSFORM_IDENTITY, NULL);
    png_write_end(png_ptr, info_ptr);
    fclose(fp);

    png_destroy_write_struct(&png_ptr, &info_ptr);
}

int main()
{
    DcmFileFormat format;
    auto status = format.loadFile("../00000300.dcm");
    if (!status.good())
    {
        std::cerr << "succeed!" << std::endl;
        return -1;
    }

    auto pdataset = format.getDataset();

    Uint16 rows;
    pdataset->findAndGetUint16(DCM_Rows, rows);
    std::cout << rows << std::endl;

    Uint16 sample_per_pixel = 0;
    if (pdataset->findAndGetUint16(DCM_SamplesPerPixel, sample_per_pixel).bad())
    {
        std::cerr << "Failed to get DCM_SamplesPerPixel" << std::endl;
        return -1;
    }

    if (sample_per_pixel != 1)
    {
        std::cerr << "sample per pixel is not 1" << std::endl;
        return -1;
    }

    Uint16 bits_allocated;
    if (pdataset->findAndGetUint16(DCM_BitsAllocated, bits_allocated).bad())
    {
        std::cerr << "Failed to get DCM_BitsAllocated" << std::endl;
        return -1;
    }

    if (bits_allocated == 8)
    {
        const Uint8 *val;
        unsigned long len = 0;
        if (pdataset->findAndGetUint8Array(DCM_PixelData, val, &len).bad())
        {
            std::cerr << "Failed to DCM_PixelData" << std::endl;
            return -1;
        }

        save_png_gray8(const_cast<uint8_t *>(val), rows);
    }

    if (bits_allocated == 16)
    {
        const Uint16 *val;
        unsigned long len = 0;
        if (pdataset->findAndGetUint16Array(DCM_PixelData, val, &len).bad())
        {
            std::cerr << "Failed to DCM_PixelData" << std::endl;
            return -1;
        }

        save_png_gray16(const_cast<uint16_t *>(val), rows);
    }

    return 0;
}